from .YouTubeDL import NODE_CLASS_MAPPINGS
import os
import nodes
from aiohttp import web
from pathlib import Path

js_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "js")

nodes.EXTENSION_WEB_DIRS["ComfyUI-yt_dl"] = js_dir

__all__ = ['NODE_CLASS_MAPPINGS']

from server import PromptServer

routes = PromptServer.instance.routes


def _get_output_dir():
    try:
        import folder_paths
        base_output = None
        if hasattr(folder_paths, 'get_output_directory'):
            base_output = folder_paths.get_output_directory()
        elif hasattr(folder_paths, 'get_temp_directory'):
            base_output = folder_paths.get_temp_directory()
        else:
            base_output = os.path.abspath(os.path.join(os.getcwd(), 'output'))
        yt_dir = os.path.join(base_output, 'yt_dl')
        os.makedirs(yt_dir, exist_ok=True)
        return yt_dir
    except Exception:
        return os.path.abspath(os.path.join(os.getcwd(), 'output', 'yt_dl'))


@routes.post('/yt_dl/api/ytdl')
async def ytdl_endpoint(request):
    try:
        import yt_dlp

        data = await request.json()
        url = data.get('url', '')
        mode = data.get('mode', 'video')
        quality = data.get('quality', 'highest')
        if not url:
            return web.json_response({'error': 'Missing url'}, status=400)

        output_dir = _get_output_dir()
        os.makedirs(output_dir, exist_ok=True)

        if mode == 'video':
            format_string = 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best'
            if quality != 'highest':
                format_string += f"[height<={quality}]"
        elif mode == 'audio':
            format_string = 'bestaudio[ext=m4a]/best[ext=m4a]'
        else:
            format_string = None

        outtmpl = os.path.join(output_dir, '%(title)s.%(ext)s')
        ydl_opts = {
            'format': format_string or 'best',
            'outtmpl': outtmpl,
            'writethumbnail': True,
            'quiet': False,
        }

        paths = []

        class CaptureHook(yt_dlp.postprocessor.common.PostProcessor):
            def run(self, information):
                filepath = information.get('filepath') or information.get('requested_downloads', [{}])[0].get('filepath')
                if filepath:
                    paths.append(filepath)
                return [], information

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            try:
                ydl.add_post_processor(CaptureHook(ydl))
            except Exception:
                pass
            ydl.download([url])

        file_path = paths[-1] if paths else ''
        return web.json_response({'file_path': file_path})
    except Exception as e:
        return web.json_response({'error': str(e)}, status=500)

